#!/bin/bash

# FLUX 24G显存优化训练脚本
# 适用于24GB VRAM的GPU - 高级优化版本

echo "========================================"
echo "FLUX 24G显存高级优化训练脚本"
echo "========================================"

# 设置环境变量
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export CUDA_LAUNCH_BLOCKING=1
export TOKENIZERS_PARALLELISM=false

# 基础路径定义
WORKSPACE="/workspace"
AI_DIR="$WORKSPACE/AI"
PYTHON_PATH="$AI_DIR/python/bin/python"
SCRIPT_PATH="$AI_DIR/ComfyUI/custom_nodes/comfyui_lora_train/sd_scripts/flux_train_network.py"

# 模型路径定义
MODELS_DIR="$AI_DIR/ComfyUI/models"
FLUX_MODEL="$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors"
CLIP_L_MODEL="$MODELS_DIR/clip/clip_l.safetensors"
T5XXL_MODEL="$MODELS_DIR/clip/t5xxl_fp8_e4m3fn.safetensors"
AE_MODEL="$MODELS_DIR/vae/ae.safetensors"

# 数据集和输出路径
TRAIN_DATA_DIR="$WORKSPACE/training_data_all/images"
OUTPUT_DIR="$WORKSPACE/training_data_all/output"
SAMPLE_PROMPTS="$WORKSPACE/training_data_all/sample_prompts.txt"

# 图片描述文件扩展名
CAPTION_EXTENSION=".txt"

# 检查必要文件是否存在
echo "检查必要文件..."

if [ ! -f "$FLUX_MODEL" ]; then
    echo "错误: FLUX模型文件不存在: $FLUX_MODEL"
    echo "请确保flux1-dev-fp8.safetensors文件在$MODELS_DIR/checkpoints/目录"
    exit 1
fi

if [ ! -f "$CLIP_L_MODEL" ]; then
    echo "错误: CLIP-L模型文件不存在: $CLIP_L_MODEL"
    echo "请确保clip_l.safetensors文件在$MODELS_DIR/clip/目录"
    exit 1
fi

if [ ! -f "$T5XXL_MODEL" ]; then
    echo "错误: T5XXL模型文件不存在: $T5XXL_MODEL"
    echo "请确保t5xxl_fp8_e4m3fn.safetensors文件在$MODELS_DIR/clip/目录"
    exit 1
fi

if [ ! -f "$AE_MODEL" ]; then
    echo "错误: AE模型文件不存在: $AE_MODEL"
    echo "请确保ae.safetensors文件在$MODELS_DIR/vae/目录"
    exit 1
fi

if [ ! -d "$TRAIN_DATA_DIR" ]; then
    echo "错误: 训练数据目录不存在: $TRAIN_DATA_DIR"
    echo "请确保training_data_all/images目录存在并包含训练图像"
    exit 1
fi

if [ ! -f "$SAMPLE_PROMPTS" ]; then
    echo "错误: 采样提示词文件不存在: $SAMPLE_PROMPTS"
    echo "请确保sample_prompts.txt文件在$WORKSPACE/training_data_all/目录"
    exit 1
fi

# 检查Python环境
echo "检查Python环境..."
if [ ! -f "$PYTHON_PATH" ]; then
    echo "错误: Python路径不存在: $PYTHON_PATH"
    exit 1
fi

# 检查accelerate
echo "检查accelerate..."
"$PYTHON_PATH" -c "import accelerate" &> /dev/null
if [ $? -ne 0 ]; then
    echo "错误: accelerate未安装"
    echo "请运行: $PYTHON_PATH -m pip install accelerate"
    exit 1
fi

# 检查GPU
echo "检查GPU..."
if ! command -v nvidia-smi &> /dev/null; then
    echo "错误: nvidia-smi未找到，请确保NVIDIA驱动已安装"
    exit 1
fi

# 检查GPU显存
echo "检查GPU显存..."
GPU_MEMORY=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1)
if [ "$GPU_MEMORY" -lt 20000 ]; then
    echo "警告: GPU显存不足24GB (当前: ${GPU_MEMORY}MB)"
    echo "建议使用24GB或更大显存的GPU"
    read -p "是否继续训练? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        echo "训练已取消"
        exit 1
    fi
fi

# 创建输出目录
if [ ! -d "$OUTPUT_DIR" ]; then
    echo "创建输出目录: $OUTPUT_DIR"
    mkdir -p "$OUTPUT_DIR"
fi

# 显示训练配置
echo
echo "========================================"
echo "训练配置信息"
echo "========================================"
echo "模型: $FLUX_MODEL"
echo "训练数据: $TRAIN_DATA_DIR"
echo "输出目录: $OUTPUT_DIR"
echo "采样提示词: $SAMPLE_PROMPTS"
echo "GPU显存: ${GPU_MEMORY}MB"
echo

# 显示训练前GPU显存使用情况
echo "训练前GPU显存使用情况："
nvidia-smi --query-gpu=memory.total,memory.used,memory.free --format=csv,noheader,nounits

# 开始训练
echo
echo "========================================"
echo "开始FLUX高级优化训练"
echo "========================================"

# 记录开始时间
START_TIME=$(date +%s)

# 设置环境变量确保使用正确的Python
export PATH="$AI_DIR/python/bin:$PATH"

# 高级优化训练命令
accelerate launch --mixed_precision bf16 --num_cpu_threads_per_process 1 --num_processes 1 --num_machines 1 --dynamo_backend no "$SCRIPT_PATH" \
--pretrained_model_name_or_path "$FLUX_MODEL" \
--clip_l "$CLIP_L_MODEL" \
--t5xxl "$T5XXL_MODEL" \
--ae "$AE_MODEL" \
--train_data_dir "$TRAIN_DATA_DIR" \
--cache_latents_to_disk \
--save_model_as safetensors \
--sdpa \
--persistent_data_loader_workers \
--max_data_loader_n_workers 2 \
--seed 42 \
--gradient_checkpointing \
--mixed_precision bf16 \
--save_precision bf16 \
--network_module networks.lora_flux \
--network_dim 8 \
--network_alpha 4 \
--network_train_unet_only \
--optimizer_type adafactor \
--optimizer_args "relative_step=False" "scale_parameter=False" "warmup_init=False" \
--learning_rate 1e-4 \
--lr_scheduler constant_with_warmup \
--max_grad_norm 0.0 \
--cache_text_encoder_outputs \
--cache_text_encoder_outputs_to_disk \

--highvram \
--max_train_epochs 4 \
--save_every_n_epochs 1 \
--output_dir "$OUTPUT_DIR" \
--output_name flux-lora-24g-optimized \
--timestep_sampling shift \
--discrete_flow_shift 3.1582 \
--model_prediction_type raw \
--guidance_scale 1.0 \
--network_args "split_qkv=True" "verbose=True" \
--sample_every_n_epochs 1 \
--sample_prompts "$SAMPLE_PROMPTS" \
--sample_sampler ddim

# 检查训练是否成功
TRAIN_EXIT_CODE=$?
if [ $TRAIN_EXIT_CODE -ne 0 ]; then
    echo "错误: 训练失败，退出代码: $TRAIN_EXIT_CODE"
    exit $TRAIN_EXIT_CODE
fi

# 计算训练时间
END_TIME=$(date +%s)
TRAIN_DURATION=$((END_TIME - START_TIME))
TRAIN_HOURS=$((TRAIN_DURATION / 3600))
TRAIN_MINUTES=$(((TRAIN_DURATION % 3600) / 60))
TRAIN_SECONDS=$((TRAIN_DURATION % 60))

echo
echo "========================================"
echo "FLUX高级优化训练完成！"
echo "========================================"
echo "训练耗时: ${TRAIN_HOURS}小时 ${TRAIN_MINUTES}分钟 ${TRAIN_SECONDS}秒"

# 显示最终显存使用情况
echo
echo "训练后GPU显存使用情况："
nvidia-smi --query-gpu=memory.total,memory.used,memory.free --format=csv,noheader,nounits

# 检查输出文件
echo
echo "检查输出文件..."
if [ -d "$OUTPUT_DIR/flux-lora-24g-optimized" ]; then
    echo "✓ 模型输出目录已创建: $OUTPUT_DIR/flux-lora-24g-optimized"
    ls -la "$OUTPUT_DIR/flux-lora-24g-optimized/"
    
    # 检查样本图像
    if [ -d "$OUTPUT_DIR/flux-lora-24g-optimized/sample_images" ]; then
        echo "✓ 样本图像目录已创建"
        SAMPLE_COUNT=$(find "$OUTPUT_DIR/flux-lora-24g-optimized/sample_images" -name "*.png" | wc -l)
        echo "样本图像数量: $SAMPLE_COUNT"
    fi
    
    # 检查训练日志
    if [ -f "$OUTPUT_DIR/flux-lora-24g-optimized/train_log.txt" ]; then
        echo "✓ 训练日志已生成"
        echo "日志文件大小: $(du -h "$OUTPUT_DIR/flux-lora-24g-optimized/train_log.txt" | cut -f1)"
    fi
else
    echo "警告: 模型输出目录未找到"
fi

echo
echo "训练完成！模型文件保存在: $OUTPUT_DIR/flux-lora-24g-optimized/"
echo "您可以在ComfyUI中加载生成的LoRA模型进行测试"




